#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/wait.h> // Added for waitpid()
#include <signal.h>   // Added for signal handling

#define PORT 8888
#define BACKLOG 5
#define BUF_SIZE 200

int conn_amount = 0; // Global counter (approximate in parent)

// Signal handler for SIGCHLD to reap zombie child processes
void sigchld_handler(int sig)
{
    while (waitpid(-1, NULL, WNOHANG) > 0)
        ;
}

// Function to handle communication with a single client (in child)
void handle_conn(int connsocket)
{
    char buf[BUF_SIZE];
    ssize_t ret;

    printf("Child pid %d handling connection on socket %d\n", getpid(), connsocket);

    while (1)
    {
        memset(buf, 0, BUF_SIZE);
        ret = recv(connsocket, buf, sizeof(buf) - 1, 0);

        if (ret < 0)
        {
            if (errno == EINTR)
                continue;
            perror("recv error in child");
            break;
        }
        else if (ret == 0)
        {
            printf("Client on socket %d closed connection (pid %d).\n", connsocket, getpid());
            break;
        }
        else
        {
            buf[ret] = '\0';
            printf("Pid %d received from socket %d: %s\n", getpid(), connsocket, buf);
        }
    }
    close(connsocket);
    printf("Child pid %d finished handling socket %d.\n", getpid(), connsocket);
    exit(0); // Terminate child process
}

int main(void)
{
    int sock_fd, newconn_fd;
    struct sockaddr_in server_addr;
    int yes = 1;
    pid_t pid;

    // Setup SIGCHLD Handler
    struct sigaction sa;
    sa.sa_handler = sigchld_handler;
    sigemptyset(&sa.sa_mask);
    sa.sa_flags = SA_RESTART;
    if (sigaction(SIGCHLD, &sa, NULL) == -1)
    {
        perror("sigaction SIGCHLD failed");
        exit(1);
    }

    // Server Setup
    if ((sock_fd = socket(PF_INET, SOCK_STREAM, 0)) == -1)
    {
        perror("socket failed");
        exit(1);
    }
    if (setsockopt(sock_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(int)) == -1)
    {
        perror("setsockopt failed");
        close(sock_fd);
        exit(1);
    }
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    server_addr.sin_addr.s_addr = INADDR_ANY;
    memset(server_addr.sin_zero, '\0', sizeof(server_addr.sin_zero));
    if (bind(sock_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) == -1)
    {
        perror("bind error");
        close(sock_fd);
        exit(1);
    }
    if (listen(sock_fd, BACKLOG) == -1)
    {
        perror("listen error");
        close(sock_fd);
        exit(1);
    }

    printf("Parent server (pid %d) listening on port %d\n", getpid(), PORT);
    conn_amount = 0;

    // Main Accept Loop
    while (1)
    {
        newconn_fd = accept(sock_fd, NULL, NULL); // Simplified accept

        if (newconn_fd < 0)
        {
            if (errno == EINTR)
            {
                printf("accept() interrupted, continuing...\n");
                continue;
            }
            else
            {
                perror("accept error");
                continue;
            }
        }

        printf("Parent pid %d accepted connection on socket %d\n", getpid(), newconn_fd);
        conn_amount++;
        printf("Parent pid %d: Approximate connection count: %d\n", getpid(), conn_amount);

        pid = fork();

        if (pid < 0)
        {
            perror("fork failed");
            close(newconn_fd); // Close accepted socket in parent
            continue;
        }
        else if (pid == 0)
        {                            // Child Process
            close(sock_fd);          // Child doesn't need listener
            handle_conn(newconn_fd); // Handle connection and exit
        }
        else
        {                      // Parent Process
            close(newconn_fd); // Parent doesn't need connected socket
            printf("Parent pid %d created child pid %d to handle socket %d\n", getpid(), pid, newconn_fd);
            // Parent continues loop
        }
    } // End while(1)

    close(sock_fd);
    return 0;
}